from collections import OrderedDict
from tqdm import tqdm
import argparse
import torch
from dataset.cad_dataset import get_dataloader, normalize_data
from config import ConfigAE
from utils import cycle
from trainer import TrainerAE
from cadlib.macro import *
from cadlib.cad_transfer import vec2sv
import torch.nn.functional as F
import wandb

wandb.init(project='deep_init')

def main():
    
    # create experiment cfg containing all hyperparameters
    cfg = ConfigAE('train')

    # create network and training agent
    tr_agent = TrainerAE(cfg)

    # load from checkpoint if provided
    if cfg.cont:
        tr_agent.load_ckpt(cfg.ckpt)

    # create dataloader
    train_loader = get_dataloader('train', cfg)
    # temp = next(iter(train_loader))

    val_loader = get_dataloader('validation', cfg)
    val_loader_all = get_dataloader('validation', cfg)
    val_loader = cycle(val_loader)

    # Set the device
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    # start training
    clock = tr_agent.clock

    for e in range(clock.epoch, cfg.nr_epochs):
        # begin iteration
        pbar = tqdm(train_loader)
        for b, data in enumerate(pbar):
            # Debugging line to check the data structure
            #print(type(data), data)  # Check the type and content of data

            # Convert each tensor in the data dictionary to the device
            data = {k: v.to(device) for k, v in data.items() if isinstance(v, torch.Tensor)}

            # train step
            outputs, losses = tr_agent.train_func(data)
            batch_out_vec = tr_agent.logits2vec(outputs)

            pbar.set_description("EPOCH[{}][{}]".format(e, b))
            pbar.set_postfix(OrderedDict({k: v.item() for k, v in losses.items()}))
            wandb.log({"train_loss_cmd": losses["loss_cmd"].item(), "train_loss_args": losses["loss_args"].item()})

            # validation step
            if clock.step % cfg.val_frequency == 0:
                val_data = next(val_loader)
                val_data = {k: v.to(device) for k, v in val_data.items() if isinstance(v, torch.Tensor)}
                commands = val_data["command"]
                conditions = val_data["conditions"]
                
                output_sv = []
                valid_count = 0
                
                with torch.no_grad():
                    outputs, losses = tr_agent.val_func(val_data)
                    batch_out_vec = tr_agent.logits2vec(outputs)
                    # temp = batch_out_vec.shape
                    # print(temp)
                    for idx in range(batch_out_vec.shape[0]):
                        seq_len = commands[idx].tolist().index(EOS_IDX)
                        out_vec = batch_out_vec[idx][:seq_len]
                        try:
                            area, vol = vec2sv(out_vec, is_mat=True)
                        except:
                            area, vol = -1, -1
                        if not area == vol == -1:
                            valid_count += 1
                            sv_data = {"area": area, "vol": vol}
                            sv_data = normalize_data(sv_data)
                            area, vol = sv_data["area"], sv_data["vol"]
                        output_sv.append([area, vol]) 
                    output_sv = torch.tensor(output_sv, dtype=torch.float32).to(conditions.device)
                    
                if valid_count == 0:
                    mse = 0.0 # 如果是，设置 MSE 为 0
                else:
                    mask = output_sv != -1
                    output_sv = output_sv[mask]
                    conditions = conditions[mask]                
                    
                    mse = F.mse_loss(conditions, output_sv).item()    
                    
                valid_pi = valid_count / batch_out_vec.shape[0]   
                wandb.log({"val_loss_cmd": losses["loss_cmd"].item(), "val_loss_args": losses["loss_args"].item(),
                           "val_mse": mse, "valid_pi": valid_pi})
                

            clock.tick()

            tr_agent.update_learning_rate()

        if clock.epoch % 5 == 0:
            tr_agent.evaluate(val_loader_all)

        clock.tock()

        if clock.epoch % cfg.save_frequency == 0:
            tr_agent.save_ckpt()

        tr_agent.save_ckpt('latest')
        
        


if __name__ == '__main__':
    main()
    
    wandb.finish()

